﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace LPLib
{
    public class Trefoil
    {
        private int[][] index;
        private Polymer<double[]> rootNormalized;
        private Polymer<double[]> mrcaNormalized;
        private List<Polymer<double[]>> polymers;
        private Polymer<double[]> rootLikelihoodSequence;
        public Polymer<double[]> RootLikelihoodSequence
        {
            get { return rootLikelihoodSequence; }
            set { rootLikelihoodSequence = value; }
        }
        private Polymer<double[]> mrcaLikelihoodSequence;
        public Polymer<double[]> MRCALikelihoodSequence
        {
            get { return mrcaLikelihoodSequence; }
            set { mrcaLikelihoodSequence = value; }
        }
        private double[] distancesToMRCA;
        private double distanceToRoot;
        public int AlignmentLength
        {
            get { return index.Length; }
        }

        private int[][] segmentStart;
        public int[][] SegmentStart
        {
            get { return segmentStart; }
        }

        private int[] segmentLength;
        public int[] SegmentLength
        {
            get { return segmentLength; }
        }

        private int nSegments;
        public int NSegments
        {
            get { return nSegments; }
        }

        public int[][] Index
        {
            get { return index; }
            set { index = value; }
        }

        private double nearlyOne;
        public List<Polymer<double[]>> Polymers
        {
            get { return polymers; }
        }

        public Polymer<double[]> RootNormalized
        {
            get { return rootNormalized; }
            set { rootNormalized = value; }
        }

        public Polymer<double[]> MRCANormalized
        {
            get { return mrcaNormalized; }
            set { mrcaNormalized = value; }
        }

        public double[] DistancesToMRCA
        {
            get { return distancesToMRCA; }
            set { distancesToMRCA = value; }
        }

        public double DistanceToRoot
        {
            get { return distanceToRoot; }
            set { distanceToRoot = value; }
        }

        public Dictionary<GeneAnnotations, int> Features;

        private double indelPartialLikelihood;
        public double IndelPartialLikelihood
        {
            get { return indelPartialLikelihood; }
            set { indelPartialLikelihood = value; }
        }

        public Trefoil(PolymerPair<double[], double[]> p1, PolymerPair<double[], double[]> p2)
        {
            nearlyOne = 1 - 0.0001;
            // merge alignment indices
            //if (p1.Polymer1.Seq.Length == p2.Polymer1.Seq.Length)
            //{
            //    index = MergeIndices(p1.Index, p2.Index);
            //}
            //else
            //{
                Tuple<int[][], double[][]> tuple = AlignIndices(p1, p2);
                index = tuple.Item1;
                rootLikelihoodSequence = new Polymer<double[]>(new NucleotidePMF(), "Root", tuple.Item2);
                rootNormalized = Polymer<byte>.Normalize(rootLikelihoodSequence);
            //}

            // create the segment list
            AlignmentToSegments();

            polymers = new List<Polymer<double[]>>();
            polymers.Add(p1.Polymer0);
            polymers.Add(p2.Polymer0);

            // transfer essential features
            Features = new Dictionary<GeneAnnotations, int>();

            // TODO: decide what to do if these two features are absent or if they are inconsistent between the two polymers.
            for (int k = 0; k < index.Length; k++)
            {
                if (index[k][0] < 0)
                    continue;

                if (p1.Features.ContainsKey(GeneAnnotations.CDR3Start) && p1.Index[p1.Features[GeneAnnotations.CDR3Start]][0] == index[k][0])
                {
                    Features.Add(GeneAnnotations.CDR3Start, k);
                }
                if (p1.Features.ContainsKey(GeneAnnotations.JInvStart) && p1.Index[p1.Features[GeneAnnotations.JInvStart]][0] == index[k][0])
                {
                    Features.Add(GeneAnnotations.JInvStart, k);
                }
            }

            DistancesToMRCA = new double[2];

        }

        public static int[][] MergeIndices(int[][] index1, int[][] index2)
        {
            List<int[]> index = new List<int[]>();
            int k1 = 0;
            int k2 = 0;

            // make sure the common reference index points to the same position in both reference sequences
            while (index1[k1][1] != index2[k2][1])
            {
                if (k1 >= index1.Length || k2 >= index2.Length)
                {
                    return null;
                }

                if (index1[k1][1] < index2[k2][1])
                {
                    k1++;
                }
                else
                {
                    k2++;
                }
            }

            int i1, i2, ia, ir; // entries for the index array
            int ca = index1[k1][1]; // coordinate for the ancestor. Initial value is the common coordinate for the two references
            while (k1 < index1.Length && k2 < index2.Length)
            {
                if (index1[k1][1] > -1 && index2[k2][1] > -1)
                {
                    // both references are ungapped here
                    i1 = index1[k1][0];
                    i2 = index2[k2][0];
                    ir = index1[k1][1];
                    k1++;
                    k2++;
                    if (i1 > -1 || i2 > -1)
                    {
                        // at least one of the observed sequences is ungapped here. Let the ancestor be ungapped, too.
                        ia = ca;
                        ca++;
                    }
                    else
                    {
                        ia = -1;
                    }
                }
                else if (index1[k1][1] > -1) // the first reference sequence is not gapped, but the second is. Thus, the second observed sequence has an insertion
                // but the other does not. Gap the mrca, the root, the first observed sequence.
                {
                    i1 = -1;
                    ir = -1;
                    ia = -1;
                    i2 = index2[k2][0];
                    k2++;
                }
                else if (index2[k2][1] > -1) // vice-versa
                {
                    i1 = index1[k1][0];
                    i2 = -1;
                    ia = -1;
                    ir = -1;
                    k1++;
                }
                else // both reference sequences have gaps, both observed sequences have insertions. 
                {
                    i1 = index1[k1][0];
                    i2 = index2[k2][0];
                    ia = ca;
                    ir = -1;
                    ca++;
                    k1++;
                    k2++;
                }
                index.Add(new int[] { i1, i2, ia, ir });
            }
            return index.ToArray();
        }

        public static Tuple<int[][], double[][]> AlignIndices(PolymerPair<double[], double[]> p0, PolymerPair<double[], double[]> p1)
        {
            double time = 0.2;
            JCEvolutionModel model = new JCEvolutionModel(time);
            Polymer<byte> p1b = TypeConverter.ToPolynucleotideByte(p0.Polymer1);
            Polymer<byte> p2b = TypeConverter.ToPolynucleotideByte(p1.Polymer1);
            MNMNAlignmentKernelSymmetric kernel = new MNMNAlignmentKernelSymmetric(model, p1b, p2b);
            PairwiseAlignerSymmetric<byte> aligner = new PairwiseAlignerSymmetric<byte>(kernel);
            int[] start = aligner.FillScoreMatrix();
            PolymerPair<byte, byte> pair = aligner.TraceBack(start);

            //// debug
            //string output = Writer<byte>.WriteFasta(new MixedNucleotideByte(), pair);
            //System.IO.File.WriteAllText("alignIndicesPair.fasta", output);
            //// debug

            List<int[]> index = new List<int[]>();
            int[][] index0 = p0.Index;
            int[][] index1 = p1.Index;

            int k0 = 0;
            int k1 = 0;
            int kR = 0;
            int k = 0;

            int i0 = 0;
            int i1 = 0;
            int iR = 0;
            List<double[]> rootConsensus = new List<double[]>();
            double[] pmf = new double[5];
            while (k0 < index0.Length && k1 < index1.Length && kR < pair.Index.Length)
            {
                k++;
                if (pair.Index[kR][1] > -1 && pair.Index[kR][0] > -1) // note that these indices cannot both be gaps
                {
                    if (index0[k0][1] > -1)
                    {
                        if (index1[k1][1] > -1) // A: both references are ungapped
                        {
                            rootConsensus.Add(consensus(p0.Polymer1.Seq[index0[k0][1]], p1.Polymer1.Seq[index1[k1][1]]));
                            i0 = index0[k0][0];
                            i1 = index1[k1][0];
                            iR = kR;
                            k0++;
                            k1++;
                            kR++;
                        }
                        else // B1: reference 0 is ungapped, reference 1 is gapped;
                        {
                            // only i1 advances
                            i0 = -1;
                            i1 = index1[k1][0];
                            k1++;
                            iR = -1;
                        }
                    }
                    else // reference 0 is gapped
                    {
                        if (index1[k1][1] > -1) // B0: reference 1 is ungapped, reference 0 is gapped
                        {
                            i0 = index0[k0][0];
                            k0++;
                            i1 = -1;
                            iR = -1;
                        }
                        else // C: both references are gapped
                        {
                            i0 = index0[k0][0];
                            k0++;
                            i1 = index1[k1][0];
                            k1++;
                            iR = -1;
                        }
                    }
                }
                else if (pair.Index[kR][0] < 0 && pair.Index[kR][1] > -1)/* one of the reference sequences has a gap in their mutual alignment
                       One must choose the reference sequence appropriately in this case.  */
                {
                    if (index0[k0][1] > -1)
                    {
                        if (index1[k1][1] > -1) // D0
                        {
                            rootConsensus.Add(p1.Polymer1.Seq[index1[k1][1]]);
                            i1 = index1[k1][0];
                            k1++;
                            i0 = -1;
                            iR = kR;
                            kR++;
                        }
                        else // E10
                        {
                            i0 = -1;
                            i1 = index1[k1][0];
                            k1++;
                            iR = -1;
                        }
                    }
                    else // index0[k0][1] < 0
                    {
                        if (index1[k1][1] > -1) // E00
                        {
                            rootConsensus.Add(p1.Polymer1.Seq[index1[k1][1]]);
                            i0 = index0[k0][0];
                            k0++;
                            i1 = index1[k1][0];
                            k1++;
                            iR = kR;
                            kR++;
                        }
                        else // F0
                        {
                            i0 = -1;
                            i1 = index1[k][0];
                            k1++;
                            iR = -1;
                        }
                    }
                }
                else // pair.Index[k][0] > -1 & pair.Index[k][1] < 0
                {
                    if (index0[k0][1] > -1)
                    {
                        if (index1[k1][1] > -1) // D1
                        {
                            rootConsensus.Add(p0.Polymer1.Seq[index0[k0][1]]);
                            i0 = index0[k0][0];
                            k0++;
                            i1 = -1;
                            iR = kR;
                            kR++;
                        }
                        else // E11
                        {
                            rootConsensus.Add(p0.Polymer1.Seq[index0[k0][1]]);
                            i0 = index0[k0][0];
                            k0++;
                            i1 = index1[k1][0];
                            k1++;
                            iR = kR;
                            kR++;
                        }
                    }
                    else // index0[k0][1] < 0
                    {
                        if (index1[k1][1] > -1) // E01
                        {
                            i0 = index0[k0][0];
                            k0++;
                            i1 = -1;
                            iR = -1;
                        }
                        else // F1
                        {
                            i0 = index0[k0][0];
                            k0++;
                            i1 = -1;
                            iR = -1;
                        }
                    }
                }
                index.Add(new int[] { i0, i1, -1, iR });
            }

            // construct MRCA index
            int ka = 0;
            for (k = 0; k < index.Count; k++ )
            {
                int count = (index[k][0] > -1 ? 1 : 0) + (index[k][1] > -1 ? 1 : 0) + (index[k][3] > -1 ? 1 : 0);
                if (count > 1)
                {
                    index[k][2] = ka;
                    ka++;
                }
            }
            return new Tuple<int[][], double[][]>(index.ToArray(), rootConsensus.ToArray());                
        }

        public void TransferNucleotides(NucleotideEvolutionModel model, int cdr3Start = 0)
        {
            double sum = 0;
            for (int k = cdr3Start; k < index.Length; k++)
            {
                if (index[k][3] > -1 && index[k][2] > -1)
                {
                    sum = rootNormalized.Seq[index[k][3]].Sum();
                    if (sum < nearlyOne)
                    {
                        double[] pmf = model.DestinationPMF(mrcaNormalized.Seq[index[k][2]], distanceToRoot);
                        double weight = 1 - sum;
                        for (int j = 0; j < 5; j++)
                        {
                            rootNormalized.Seq[index[k][3]][j] += weight * pmf[j];
                        }
                    }
                }
            }
        }

        public void AlignmentToSegments()
        {
            List<int[]> segmentStartAsList = new List<int[]>();
            segmentStartAsList.Add(new int[4]);
            List<int> segmentLengthAsList = new List<int>();
            segmentLengthAsList.Add(0);

            int segment = 0;
            for (int i = 0; i < 4; i++ )
            {
                segmentStartAsList[0][i] = Index[0][i];
            }

            for (int k = 1; k < Index.Length - 1; k++)
            {
                if (isBreakpoint(k))
                {
                    segment++;
                    segmentStartAsList.Add(new int[4]);
                    segmentLengthAsList.Add(0);

                    for (int i = 0; i < 4; i++ )
                    {
                        if (Index[k + 1][i] < 0) // the new segment is absent in sequence 0...
                        {
                            segmentStartAsList[segment][i] = -1;
                        }
                        else
                        {
                            segmentStartAsList[segment][i] = Index[k + 1][i];
                        }
                    }

                    // compute length
                    for (int i = 0; i < 4; i++ )
                    {
                        if (segmentStartAsList[segment - 1][i] > -1)
                        {
                            segmentLengthAsList[segment - 1] = Index[k][i] + 1 - segmentStartAsList[segment - 1][i];
                            break;
                        }
                    }
                }
            }

            // finish up
            nSegments = segment + 1;

            for (int i = 0; i < 4; i++)
            {
                if (segmentStartAsList[segment][i] > -1)
                {
                    segmentLengthAsList[segment] = Index.Last()[i] + 1 - segmentStartAsList[segment][i];
                    break;
                }
            }

            segmentLength = segmentLengthAsList.ToArray();
            segmentStart = segmentStartAsList.ToArray();
        }

        private static double[] consensus(double[] p1, double[] p2)
        {
            if (p1.Length != p2.Length)
                throw new ArgumentException("The pmfs have different dimensions.");

            double[] consens = new double[p1.Length];
            double sum = 0;
            for (int i = 0; i < p1.Length; i++)
            {
                consens[i] = p1[i] * p2[i];
                sum += consens[i];
            }

            if (sum == 0)
            {
                for (int i= 0; i< consens.Length; i++)
                {
                    consens[i] = (p1[i] + p2[i]) / 2;
                }
            }
            else
            {
                for (int i = 0; i < consens.Length; i++)
                {
                    consens[i] /= sum;
                }
            }
            return consens;
        }

        private bool isBreakpoint(int k)
        {
            for (int i= 0; i < 4; i++)
            {
                if ( (Index[k][i] < 0 && Index[k + 1][i] > -1 ) || (Index[k][i] > -1 && Index[k + 1][i] < 0) )
                    return true;
            }
            return false;
        }
    }

}
